#!/bin/bash

# TERMINAL COLORS -----------------------------------------------------------------

NONE='\033[00m'
RED='\033[01;31m'
GREEN='\033[01;32m'
YELLOW='\033[01;33m'

# INTERFACES

EXT_NIC="eth0"
INT_NIC="virbr0"
IP="192.168.2.2"

TC="/sbin/tc"
IPT="/sbin/iptables"
#U32="$TC filter add dev $EXT_NIC protocol ip parent 1:0 prio 1 u32"

#  tc uses the following units when passed as a parameter.
#  kbps: Kilobytes per second 
#  mbps: Megabytes per second
#  kbit: Kilobits per second
#  mbit: Megabits per second
#  bps: Bytes per second 

# DO NOT EDIT BELOW THIS LINE ______________________________________________
totaldown=0
totalup=0
total_clients=0
total_groups=0

groups_index=()
groups_name=()
groups_id=()
groups_down=()
groups_up=()
groups_aggr=()
groups_prio=()
groups_active=()
groups_client_count=()
groups_sub_count=()

all_ip=()
all_parentid=()
all_classid=()

#
# increment and set to 2 decimal places
#
function inc_2_leadzeroes {
        local val=$1
        val=$(($val+1))

        if [[ $val -lt 10 ]]; then
                val="0"$val
                echo "$val"
                return $val
        fi
        echo "$val"
        return $val
}

#
# increment and set to 3 decimal places
#
function inc_3_leadzeroes {
        local val=$1
        val=$(($val+1))

        if [[ $val -lt 10 ]]
        then            
                val="00"$val
                echo "$val"
                return $val
        elif [[ $val -gt 9 && $val -lt 100 ]]
        then
                val="0"$val
                echo "$val"
                return $val
        else
                echo "$val"
                return $val
        fi
}

#
# Check if IP address is valid
#
function validate_IP {
        local ip=$1
	local stat=1
	# Check the IP address under test to see if it matches the extended REGEX

	if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
		# Record current default field separator
		OIFS=$IFS
		# Set default field separator to .
		IFS='.'
		# Create an array of the IP Address being tested
		ip=($ip)
		# Reset the field separator to the original saved above IFS=$OIFS
		# Is each octet between 0 and 255?
		[[ ${ip[0]} -le 255 && ${ip[1]} -le 255 && ${ip[2]} -le 255 && ${ip[3]} -le 255 ]]
		# Set the return code.
		stat=$?
	fi
}

#
# mark packets for IP with class ID
#
function ipt_classify_ip { #ip #class
        local ip=$1
        local class=$2
        $IPT -A POSTROUTING -t mangle -o $INT_NIC -d $ip -j CLASSIFY --set-class 1:1${class}
        $IPT -A POSTROUTING -t mangle -o $EXT_NIC -s $ip  -j CLASSIFY --set-class 1:2${class}
}

#
# Classify all IPs loop
#
function ipt_classify_all_ip {
                local xsize=${#all_ip[@]}
                for (( x=0; x <xsize; x++ ))
                do
                        ipt_classify_ip ${all_ip[$x]} ${all_classid[$x]} #ip #class  
                done
}

#TODO
function tc_print_counters {
        echo "total download:$totaldown upload:$totalup"
        echo "==============================================="
        echo -e "${YELLOW}total ${TOTALGROUPS} groups and ${TOTALSUBS} subgroups ${NONE}"
}

function tc_remove {
        echo "REMOVING ROOT QDISC"
        $TC qdisc del dev $INT_NIC root
        $TC qdisc del dev $EXT_NIC root
        $TC qdisc del dev $INT_NIC ingress
        $TC qdisc del dev $EXT_NIC ingress
}

function tc_add_group { # $parent, $classid, $total, $ceil
        local parent=$1
        local classid=$2
        local total=$3
        local ceil=$4
        local DEV=$5
        $TC class add dev $DEV parent $parent classid $classid htb rate ${total}kbit ceil ${ceil}kbit   
}

function tc_add_group_down {
        #echo "PRE tc_add_group_DOWN totaldown: $totaldown $1 $2 $3 $4 $EXT_NIC"
        tc_add_group $1 $2 $3 $4 $EXT_NIC
        totaldown=$(($(($totaldown))+$(($4))))
        #echo "POST tc_add_group_DOWN totaldown: $totaldown"
}

function tc_add_group_up {
        #echo "PRE tc_add_group_UP totalup: $totalup $1 $2 $3 $4 $INT_NIC"
        tc_add_group $1 $2 $3 $4 $INT_NIC
        totalup=$(($(($totalup))+$(($4))))
        #echo "POST tc_add_group_UP totalup: $totalup"
}

function tc_root {
        echo "ADDING ROOT QDISC"
        $TC qdisc add dev $INT_NIC root handle 1: htb default 1
        $TC qdisc add dev $EXT_NIC root handle 1: htb default 1

        echo "$TC class add dev $INT_NIC parent 1: classid 1:1 htb rate ${totaldown}Kbit ceil ${totaldown}Kbit buffer 1600"
        echo "$TC class add dev $EXT_NIC parent 1: classid 1:2 htb rate ${totalup}Kbit ceil ${totalup}Kbit buffer 1600  "

        $TC class add dev $INT_NIC parent 1: classid 1:1 htb rate ${totaldown}Kbit ceil ${totaldown}Kbit buffer 1600
        $TC class add dev $EXT_NIC parent 1: classid 1:2 htb rate ${totalup}Kbit ceil ${totalup}Kbit buffer 1600
}

function tc_show {
        echo "SHOW ROOT QDISC"
        $TC -s qdisc show dev $EXT_NIC
        $TC -s class show dev $EXT_NIC
        $TC -s qdisc show dev $INT_NIC
        $TC -s class show dev $INT_NIC
}

##########################################################################
# load groups from definition file and compute traffic values like ceil  #
##########################################################################
function prepare_group_definitions {
        local z=0
        #check if group.definitions exist, if not, create it with default Fiber30 group
        if [ -s "config/group.definitions" ]; then
                echo -e "${GREEN}found group.definitions and processing...${NONE}" >> /dev/null
        else
                echo -e "${RED} group.definitions NOT FOUND!!! recreating...${NONE}"
                mkdir config
                touch config/group.definitions;
                echo "#name           download        upload          aggregation     prio_group" > config/group.definitions
                echo "Fiber30         30720           5120            8               0" >> config/group.definitions
                exit
        fi

        #read group.definitions and fill arrays with values
        while read line
        do
                #take line and split it into array, first value is main! (aka group,client..etc)
                arrl=($(echo $line | tr " " "\n"))
                if [ -z "$line" ]; then
                        echo "EMPTY LINE"
                else 
                        case "${arrl[0]}" in
                '#'*)  
                        ;;
                *)      echo "---------------------------------------------------------------------------------------------------------"
                        echo -e "${GREEN}found GROUP called ${arrl[0]} | download $((arrl[1] / 1024)) Mbs | upload $((arrl[2]/1024)) Mbs | aggr ${arrl[3]} ${NONE}"
                      	echo "---------------------------------------------------------------------------------------------------------"
                        groups_name[$z]=${arrl[0]}
                        groups_index[$z]=$((${z}+1))
                        groups_down[$z]=${arrl[1]}
                        groups_up[$z]=${arrl[2]}
                        groups_aggr[$z]=${arrl[3]}
                        groups_prio[$z]=${arrl[4]}
                        groups_sub_count[$z]=0
                        groups_active[$z]=0 #by default mark group inactive

                        #if file .group exist, MARK GROUP ACTIVE
                        if [ -f config/${arrl[0]}.group ];
                        then
                            groups_active[$z]=1
                        else
                            echo -e "${RED}#### ERROR! #### File ${groups_name[$z]}.group does NOT EXIST but is defined in group.definitions${NONE}"
                            echo "creating dummy .group file"
                            touch config/${arrl[0]}.group
                            echo "127.0.0.0" > config/${arrl[0]}.group
                            echo -e "${RED}FIX config/${arrl[0]}.group FILE BEFORE RUNNING AGAIN!!!${NONE}"
                            exit #exit so user can change 127.0.0.1 in .group file
                        fi
                        
                        z=$((${z}+1))
                        ;;
                esac
                fi
        done <config/group.definitions

        echo "END OF DEFINITIONS"       
}

function gen_subgroups { # $groups_index[] ${groups_id[$i]}
            #cat config/${groups_name[${i}]}.group|sort -n|uniq 1>/dev/null

	local dd=0
	local cc=0
	local ddd
	local i=0
        while read line
        do
		if [[ cc -lt ${groups_aggr[$1]} ]]; then
			cc=$(($cc+1))
		else 
			dd=$(($dd+1))	
			cc=0
		fi
		ddd=`inc_2_leadzeroes $dd`
	
		all_ip[${#all_ip[*]}]=$line
		all_groupid[${#all_groupid[*]}]=$2
		all_classid[${#all_classid[*]}]=$((1$2${ddd}))
		all_classid[${#all_classid[*]}]=$((2$2${ddd}))
		#echo "all_classid[1] is $1 and ${all_classid[$1]}"
		#echo "$line | cc $cc | groupid $2 | classid $ddd | spolu 1$2$ddd | check ${all_classid[$1]} overall size ${#all_classid[*]}"
		i=$(($i+1))
	done <config/${groups_name[$1]}.group
}

#function add_to_all_ip { #$line
#	echo "add_to_all_ip $1"
#	allip_size=${#all_ip[*]}
#	all_ip[$allip_size]=$1
#
#	echo "new all ip size is $(($allip_size+1)) with $1"
#}

#################################################
# get all clients for certain group             #
#################################################
function count_clients {
        #
        # Count clients and load IPs into all_ip[]
        #
        local size=${#groups_index[@]}
        echo "counting clients of $size groups"
        for (( i=0; i <size; i++ ))
        do
            #cat config/${groups_name[${i}]}.group|sort -n|uniq 1>/dev/null
            while read line
            do
		# VALIDATE IP ADDRESS
		local templine=$line
		validate_IP $templine
		local valid=$?         
		if [[ $valid -ne 0 ]]; then			
			echo -e "${RED}IP ADDRESS $line IS NOT VALID!!!${NONE}"
			exit		
		else			
			#echo "valid IP $line" 
			allip_size=${#all_ip[*]}
			all_ip[$allip_size]=$line
		fi
                
                total_clients=$((${total_clients}+1))
                groups_client_count[$i]=$((${groups_client_count[$i]}+1))

                #add_to_all_ip $line

            done <config/${groups_name[$i]}.group
        done
        echo "FOUND $total_clients CLIENTS IN TOTAL"
        
        #
        # Compute number of subgroups depending on aggregations
        #
        for (( i=0; i <size; i++ ))
        do
	        if [[ ${groups_client_count[$i]} -lt ${groups_aggr[$i]} ]];then
 	               groups_sub_count[$i]=1
                else
			groups_sub_count[$i]=$(($(($((${groups_client_count[$i]})) / $((${groups_aggr[$i]}))))+1))
                fi
                
                echo "got ${groups_name[$i]} with ${groups_client_count[$i]} clients divided into ${groups_sub_count[$i]} subs"
                totaldown=$(($((${totaldown}))+$((${groups_down[$i]}))))
                totalup=$(($((${totalup}))+$((${groups_up[$i]})) ))
        done

        #
        # Generate class ID for each group, subgroup, IP..
        #
        for (( i=0; i <size; i++ ))
        do
                local grid=`inc_2_leadzeroes $i`
                #echo "gen 1:1${grid} and 1:2${grid}"
		groups_id[$i]="${grid}"
                if [[ ${groups_active[$i]} == 1 ]]; then
			#echo "active"
			gen_subgroups $i ${groups_id[$i]}
		fi
	done
}

#FIX
function print_ip_with_classid { #1:1${sgrid}
	local size=${#all_ip[*]}

	#echo "| looking for $1"
	#echo "looping $size "
	for (( m=0; m<$size; m++ )) 
	do
		#echo "checking ${all_ip[$i]} :  $1 compare ${all_classid[$i]}"

		if [[ ${all_classid[$m]} == $1 ]];
		then
			echo "  |       |                   |     - ${all_ip[$m]} ${all_classid[$m]}" #>> /dev/null
		else
			echo "no match: ${all_classid[$m]}" >> /dev/null
		fi
	done
}

#
# Generate only tree visualization of leaves, don't create anything
#
function tc_generate_fake_leaves {
        local size=${#groups_index[*]}
        
        echo -e "${GREEN}==============================${NONE}"
        echo -e "${GREEN}= Shaper tree visualization$ =${NONE}"
        echo -e "${GREEN}==============================${NONE}"

	echo "  |"
	echo "[root qdisc]"
	echo "  |"
	echo "  1:1------"
	echo "  |       |"

	for (( i=0; i<$size; i++ ))
        do
		pgrid=`inc_2_leadzeroes $i`
		echo "  |       |"
		echo "  |       |-[${groups_name[$i]}]-1${pgrid}---------"
		echo "  |       |                   |"	
		
		#process sub groups
		subsize=${groups_sub_count[$i]}
		for (( z=0; z<$subsize; z++ ))
		do
			sgrid=`inc_2_leadzeroes $z`
			echo "  |       |                   |"	
			echo "  |       |                   |--sub 1${pgrid}${sgrid}-"

			print_ip_with_classid $((1${pgrid}${sgrid}))
		done		
	done

	echo "  |"
	echo "  1:2------"
	echo "  |       |"



        #for (( i=0; i<$size; i++ ))
        #do
		#siz=${groups_sub_count[$i]}
                #echo "|"                
                #if [[ ${groups_active[$i]} == 1 ]]; then
			
			
			#echo "| parent 1:1${groups_id[$i]} as download with $siz sub-groups"
			#for (( j=0; j<$siz; j++ ))
                    	#do
                        	#sgrid=`inc_2_leadzeroes $j`
                        	#echo "|____ 1:1${groups_id[$i]} class 1:1${groups_id[$i]}${sgrid}"
                            
                        	#print_ip_with_classid 1${groups_id[$i]}${sgrid}
                    	#done
                
			#echo "| parent 1:2${groups_id[$i]} as upload with $siz sub-groups"
                    	#for (( jj=0; jj<$siz; jj++ ))
                    	#do
                        	#sgrid=`inc_2_leadzeroes $jj`
                            	#echo "|____ 1:2${groups_id[$i]} class 1:2${groups_id[$i]}${sgrid}"
                            
                            	#print_ip_with_classid 2${groups_id[$i]}${sgrid}
                    	#done
		#else
			#echo "group ${groups_active[$i]} not active"
                #fi
                
                #echo "|"
        #done
}

#FIX FIX FIX
function tc_generate_leaves {
        local i=0
        local j=0
        local size=${#groups_id[@]}

        for (( i=0; i <size; i++ ))
        do
                grid=`inc_2_leadzeroes $i`
                echo "| create GROUP 1:1${grid}"
                echo "|"
                tc_add_group_down 1:1 1:1${grid} ${groups_down[$i]} ${groups_down[$i]} #$parent, $classid, $total, $ceil
                tc_add_group_up 1:2 1:2${grid} ${groups_up[$i]} ${groups_up[$i]}  #$parent, $classid, $total, $ceil

                echo "..."
                for (( j=0; j<${groups_sub_count[$i]}; j++ ))
                do
                        sgrid=`inc_3_leadzeroes $j`
                        echo "|___________ 1:1${grid} class 1:1${sgrid}"
                        echo "|___________ 1:2${grid} class 1:2${sgrid}"



                        #tc_add_group_down 1:1${grid} 1:1${sgrid} DOWN UP #$parent, $classid, $total, $ceil   FIX!!!!                   
                        #tc_add_group_up 1:2${grid} 1:2${sgrid} DOWN UP #$parent, $classid, $total, $ceil                       
                done
        done
}

function tc_fake_start {
        #echo "############################################################################################"

        prepare_group_definitions
        count_clients
        tc_generate_fake_leaves
       
        #echo "############################################################################################"
}

function tc_start {
        tc_remove
        tc_show
        tc_print_counters
        echo "############################################################################################"

        prepare_group_definitions
        count_clients
        tc_root
        tc_generate_leaves
        ipt_classify_all_ip
        
        echo "############################################################################################"
        tc_print_counters
        tc_show
        tc_remove
}

#MAIN ######################################################################
        echo "------------------------------------------------------------------------------------------------------"
        echo "    _____                                                     _____                                   "
        echo " __|__   |__  ____    ______  ____    __  _____  ____   _  __|___  |__  __   _  ____    _____  ______ "
        echo "|     \     ||    \  |   ___||    \  /  |/     \|    \ | ||   ___|    ||  |_| ||    \  |     ||   ___|"
        echo "|      \    ||     \ |   ___||     \/   ||     ||     \| |    --_     ||   _  ||     \ |    _||   ___|"
        echo "|______/  __||__|\__\|______||__/\__/|__|\_____/|__/\____||______|  __||__| |_||__|\__\|___|  |______|"
        echo "   |_____|                                                   |_____|                                  "
        echo "------------------------------------------------------------------------------------------------------"
        echo "DaemonShape dynamic traffic shaper. https://code.google.com/p/bashtools/ "
        echo "------------------------------------------------------------------------------------------------------"

case "$1" in
start)
        tc_start
        ;;
printtree)
        tc_fake_start
        ;;
stop)
        tc_remove
        ;;
        
restart) tc_remove
        tc_start
        ;;

status)
        echo "=IPSET=====================";
        ipset -L;
        echo "=IPTABLES=====================";
        iptables -L -n -v --line-numbers;
        ;;

version)
        echo "==================================================";
        echo "=DaemonShape version 0.1 by peter.ducai@gmail.com=";
        echo "==================================================";
        ;;

*)
        echo "usage: $0 (start|printtree|fakestarttc|generatetc <generated.file>|stop|restart|status|version)"
        exit 1
esac

exit $?
